#!/usr/bin/env python3
import os
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed
from jinja2 import Environment, BaseLoader

from common import AuthInfo, BaseMachineInfo, OSDeploy


def get_arg(key):
    for arg in sys.argv[1:]:
        if arg.startswith(f"{key}="):
            return arg.split("=", 1)[1]
    return None


def render_user_data(template_path, driver_version, cuda_version):
    env = Environment(loader=BaseLoader())
    with open(template_path, "r", encoding="utf-8") as f:
        template = env.from_string(f.read())
    return template.render(
        driver_version=driver_version,
        cuda_version=cuda_version,
    )


def deploy_one(host, maas_url, api_key, distro_series, template_path, driver_version, cuda_version):
    rendered_user_data = render_user_data(template_path, driver_version, cuda_version)

    auth = AuthInfo(maas_url, api_key)
    deployer = OSDeploy(maas_url, auth.session, host, distro_series, rendered_user_data)

    if host.get("status") == "Ready":
        deployer.deploy()
        return host, True
    else:
        return host, False


if __name__ == "__main__":
    maas_url = os.getenv("maas_url")
    api_key = os.getenv("api_key")

    distro_series = get_arg("distro_series")
    hosts_arg = get_arg("hosts")
    driver_version = get_arg("driver_version")
    cuda_version = get_arg("cuda_version")

    if not all([maas_url, api_key, distro_series, hosts_arg, driver_version, cuda_version]):
        print("Usage: deploy_OS.py distro_series=<series> hosts=<hosts_arg> driver_version=<drv> cuda_version=<cuda>")
        sys.exit(1)

    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    template_path = os.path.join(project_root, "user-data", "user-data.yml")

    if not os.path.exists(template_path):
        print(f"❌ 未找到 user-data 模板: {template_path}")
        sys.exit(2)

    auth = AuthInfo(maas_url, api_key)
    machine_info = BaseMachineInfo(auth.session, maas_url, hosts_arg)

    with ProcessPoolExecutor(max_workers=3) as executor:
        futures = [
            executor.submit(
                deploy_one,
                host,
                maas_url,
                api_key,
                distro_series,
                template_path,
                driver_version,
                cuda_version,
            )
            for host in machine_info.results
        ]
        for f in as_completed(futures):
            hostname, success = f.result()
            if success:
                print(f"✅ 部署完成: {hostname}")
            else:
                print(f"⚠️  跳过（状态非Ready）: {hostname}")
